import json

import websocket
import ssl
from src.spark.spark_api import Ws_Param, on_error, on_close, on_open


class SparkClient:
    def __init__(self, appid, api_key, api_secret, Spark_url, domain):
        self.appid = appid
        self.api_key = api_key
        self.api_secret = api_secret
        self.Spark_url = Spark_url
        self.domain = domain
        self.answer = ""

    def get_response(self, question):
        wsParam = Ws_Param(self.appid, self.api_key, self.api_secret, self.Spark_url)
        wsUrl = wsParam.create_url()
        ws = websocket.WebSocketApp(wsUrl,
                                    on_message=self.on_message,
                                    on_error=on_error,
                                    on_close=on_close,
                                    on_open=on_open)
        ws.appid = self.appid
        ws.question = question
        ws.domain = self.domain
        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
        return self.answer

    def on_message(self, ws, message):
        data = json.loads(message)
        code = data['header']['code']
        if code != 0:
            print(f'请求错误: {code}, {data}')
            ws.close()
        else:
            choices = data["payload"]["choices"]
            status = choices["status"]
            content = choices["text"][0]["content"]
            print(content, end="")
            self.answer += content
            if status == 2:
                ws.close()
